'''
Author: devis.dong
Email: devis.dong@gmail.com
Date: 2021-12-09 21:19:45
LastEditTime: 2022-12-12 22:45:01
LastEditors: devis.dong
Description:
'''

from mynets import *


class getModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
                        PointNetEncoder(dims=[3,64,64,128,1024]),
                        nn.Linear(1024, 512), nn.BatchNorm1d(512), nn.ReLU(),
                        nn.Linear(512, 256), nn.Dropout(p=0.4), nn.BatchNorm1d(256), nn.ReLU(),
                        nn.Linear(256, 128)
                    )

    def feature(self, x: torch.Tensor):
        while x.dim() < 3:
            x = x.unsqueeze(0)
        x = self.encoder(x)
        x = torch.exp(x)
        return x # [B, k]

    def fsimilarity(self, f0, f1):
        while f0.dim() < 2:
            f0 = f0.unsqueeze(0)
        while f1.dim() < 2:
            f1 = f1.unsqueeze(0)

        # 欧式距离相似度
        # sm = torch.exp(-torch.sum((f0-f1)**2, dim=1))

        # 向量夹角（余弦距离）相似度
        sm = F.cosine_similarity(f0, f1, dim=1)

        # pearson相关系数
        # vx = f0 - torch.mean(f0, dim=1, keepdim=True).repeat(1, f0.shape[1])
        # vy = f1 - torch.mean(f1, dim=1, keepdim=True).repeat(1, f1.shape[1])
        # coef = torch.sum(vx * vy, dim=1) / (torch.sqrt(torch.sum(vx ** 2, dim=1)) * torch.sqrt(torch.sum(vy ** 2, dim=1)))
        # sm = (coef + 1) * 0.5

        return sm.squeeze()

    def forward(self, x0, x1):
        return self.fsimilarity(self.feature(x0), self.feature(x1))


# 自定义损失
class getLoss(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, sm, lb):
        sm = sm.view(-1, 1)
        softmax = torch.cat([1-sm, sm], dim=1) + 1e-10
        log_softmax = torch.log(softmax)
        loss = F.nll_loss(log_softmax, lb, size_average=False)
        return loss
